from collections import defaultdict
from queue import Queue
import copy
import random


CRATE_TASK_RATIO = 0.5
MAX_EXECUTION_TIME = 5
def setConfig(create_time_ratio, max_execution_time):
    global CRATE_TASK_RATIO, MAX_EXECUTION_TIME
    CRATE_TASK_RATIO = create_time_ratio
    MAX_EXECUTION_TIME = max_execution_time
    return

# 先写个多智能体的
class Task():
    def __init__(self, start_node_index, required_resources, time_cost, system):
        
        self.start_node_index = start_node_index
        self.start_node = system.nodes[start_node_index]
        self.execute_node = None

        self.required_resources = required_resources
        # self.end_node_id = end_node_id

        
        self.time_cost = time_cost
        assert self.time_cost > 0

        self.executed_node = None
        self.have_spent_time = 0  # 如果正在执行，每个time_frame 会加1

        self.create_time = system.now_time_frame_index
        self.start_execute_time = None #开始执行的时间
        self.start_wait_time = None

        self.system = system


        return

    def calNetworkTimeCostTo(self, target_node_index):

        required_resources = self.required_resources
        system = self.system

        assert required_resources[0] == 0 and len(required_resources) == 1
        resource_size = system.data_table[required_resources[0]]['size']

        network_time_cost = resource_size / system.network_speed[(self.start_node_index, target_node_index)]

        return network_time_cost

    @property
    def now_time_frame_index(self):
        return self.system.now_time_frame_index

    def step(self):
        self.have_spent_time += 1

        if self.start_execute_time is None:
            self.start_execute_time = self.now_time_frame_index

        return
    
    @property
    def is_finish(self):
        assert self.time_cost >= self.have_spent_time
        return self.time_cost == self.have_spent_time

    @property
    def has_transform(self):
        return self.create_time + self.network_time_cost < self.now_time_frame_index

class Node():
    def __init__(self, index, system):
        self.task_queue = Queue()  #预报要执行的任务

        self.executed_tasks = []  #已经执行完的任务

        self.parallel_number = 3  #可并行的数量
        self.executing_tasks = [None] * self.parallel_number  #正在执行的任务, 一个任务对应一个cpu
        
        self.index = index  # 节点的id

        self.system = system #
        return

    def __hash__(self):
        return hash(self.index)

    @property
    def now_time_frame_index(self):
        return self.system.now_time_frame_index
    
    # 未来可以加上如果可以自己执行就自己执行，如果不行就return出来

    # 现在假设是完全中心化的
    def genTaskAndDispatch(self):
        MAX_TASK_NUMBER = 1  # 先假设只生成一个task
        require_remote_tasks = []

        for _ in range(MAX_TASK_NUMBER):
            if random.random() > CRATE_TASK_RATIO:
                execution_time = int(random.random() * MAX_EXECUTION_TIME) + 1
                new_task = Task(self.index, [0], execution_time, self.system)  # todo: 先
                require_remote_tasks.append(new_task)

        return require_remote_tasks
    
    # 接受自己或者外部来的Task
    def receiveTask(self, task):
        task.executed_node = self
        self.task_queue.put(task)
        task.start_wait_time = self.now_time_frame_index
        return

    def assignTaskForCPU(self, cpu_index):
        new_task = None
        if not self.task_queue.empty():
            new_task = self.task_queue.get()
            self.executing_tasks[cpu_index] = new_task
        return new_task

    def step(self):
        for cpu_index, task in enumerate(self.executing_tasks):
            if task is None:
                task = self.assignTaskForCPU(cpu_index)

            if task is None:
                continue

            task.step()

            if task.is_finish:
                self.executed_tasks.append(task)

                task = None
                self.executing_tasks[cpu_index] = None
        
        
        return self.genTaskAndDispatch()

    def heartbeat(self):
        executed_tasks = self.executed_tasks[-5:]

        if len(executed_tasks) == 0:
            avg_execution_time = 0.5
        else:
            avg_execution_time = sum([task.time_cost for task in executed_tasks]) / len(executed_tasks)

        return {
            'avg_execution_time': avg_execution_time,
            'queue_size': self.task_queue.qsize(),
        }
    
class GlobaDataTable():
    def __init__(self):
        self.resource2nodes = defaultdict(list)
        return

class System():
    def __init__(self, time_frame_number = 100):
        random.seed(10)  #保证每次的随机的结果应该是一样的

        self.time_frames = list(range(time_frame_number))
        self.now_time_frame_index = 0

        self.nodes = []
        self.network_speed = {}

        self.heartbeats = []
        self.heartbeat_rate = 5   # 每20个frame收集一次

        self.data_table = {
            0: {
                'size': 1,  #网速小于1
                'nodes': [],  # 哪都没有存
            },  #这个数据哪个节点都没有
        }


        self.transform_tasks = []
        
        return


    # 生成Nge节点
    def createNodes(self, node_number):
        assert node_number > 2
        assert len(self.nodes) == 0

        for i in range(node_number):
            new_node = Node(i, self)
            self.nodes.append(new_node)

        self.createRandomNetworks()
        self.getHearBeat()
        return

    def createRandomNetworks(self):
        assert len(self.nodes) > 2

        for index1, node1 in enumerate(self.nodes):
            for index2, node2 in enumerate(self.nodes):
                if node1 == node2:
                    random_speed = 10  # 自己传自己是最快的
                else:
                    random_speed = random.random() #(index1 + index2) % 5
                # self.network_speed[(node1, node2)] = random_speed
                # self.network_speed[(node2, node1)] = random_speed
                self.network_speed[(node1.index, node2.index)] = random_speed * 5
                self.network_speed[(node2.index, node1.index)] = random_speed * 5

    def scheduleAll(self, tasks):
        return [(self.schedule(task), task) for task in tasks] 

    def step(self):
        if self.now_time_frame_index % self.heartbeat_rate == 0:
            self.getHearBeat()

        new_tasks = []
        for node in self.nodes:
            new_tasks += node.step()

        for target_node_index, task in self.scheduleAll(new_tasks):
            network_time_cost = task.calNetworkTimeCostTo(target_node_index)

            task.execute_node_index = target_node_index
            task.network_time_cost = network_time_cost

            self.transform_tasks.append(task)
    

        still_transform_task = []
        for task in self.transform_tasks:
            if task.has_transform:
                target_node_index = task.execute_node_index
                self.nodes[target_node_index].receiveTask(task)
                
            else:
                still_transform_task.append(task)

        self.transform_tasks = still_transform_task

        self.now_time_frame_index += 1

        return

    def getHearBeat(self):
        heartbeats = []
        for node in self.nodes:
            heartbeats.append(node.heartbeat())

        self.heartbeats = heartbeats

    # 默认是random的
    def schedule(self, task):
        # node_index = 0
        node_index = random.choice(self.nodes).index
        return node_index

# def dumps(system):
#     all_executed_tasks = []
#     for node in system.nodes:
#         all_executed_tasks += node.executed_tasks
        
#     task_infos = {
#         {
#             'time_cost': task.time_cost,
#             'resource_size':  1,
#             'start_node': task.start_node_index,
#         }
#         for task in all_executed_tasks
#     }

#     return {
#         'task_infos': all_executed_tasks,
#         'network_speed': system.network_speed,
#     }

def analysis(system):

    all_executed_tasks = []
    for node in system.nodes:
        all_executed_tasks += node.executed_tasks

    wait_time_costs = [task.start_execute_time - task.start_wait_time for task in all_executed_tasks]
    print('execute task number: ', len(all_executed_tasks), 'average wait time', sum(wait_time_costs)/len(wait_time_costs))
    return

if __name__ == '__main__':
    snap_shoot = []  #存了快照

    system = System()
    system.createNodes(5)
    for time_frame in system.time_frames:
        system.step()

        # snap_shoot.append(copy.deepcopy(system))

        pass

    # 分析一下等待时间这些
    # 不同schedule生成的应该得是一样的才行呀

    analysis(system)
    pass